import json
from tqdm import tqdm
import re
from count_data import iter_weibo, iter_novel, iter_tnews
from collections import Counter


def count_word(num=None):
    word__count = Counter()
    for get in [iter_novel, iter_weibo, iter_tnews]:
        dtype = str(get).split(' ')[1][5:]
        total, data = get(num)
        for txt in tqdm(data, total=total, desc=dtype):
            try:
                words = re.findall('[a-zA-Z]+', txt)
                word__count.update(words)
            except:
                pass

        path = f'data/英文_{dtype}_word__count.json'
        with open(path, 'w', encoding='utf8') as f:
            json.dump(dict(word__count), f, ensure_ascii=False, indent=2)


if __name__ == '__main__':
    count_word(None)
